package main

import (
	"flag"
	"io"
	"log"
	"net"
	"os"
	"strings"
)

var (
	server  bool
	client  bool
	addr    string
	remote  string
	initx   bool
	encrypt string
)

func init() {
	flag.BoolVar(&server, "server", false, "run as server")
	flag.BoolVar(&client, "client", false, "run as client")
	flag.BoolVar(&initx, "init", false, "generate rsa file")
	flag.StringVar(&encrypt, "encrypt", "aes", "encrypt type")
	flag.Parse()
	if initx {
		// RSAKeyGen(1024)
		AesKeyGen(16)
		os.Exit(0)
	}
	if server == true && client == true {
		log.Fatalln("only run server or client")
	}
	if len(flag.Args()) < 1 {
		log.Fatalln("address missing")
	}
	d := strings.Split(flag.Arg(0), "@")
	addr = d[0]
	remote = d[1]
}
func main() {
	loadAesKey()
	if server == true {
		// loadPrivate()
		same(addr, remote)
	}
	if client == true {
		// loadPublic()
		same(addr, remote)
	}
}
func same(local, remote string) {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatalln(err)
	}
	for {
		a, err := l.Accept()
		if err != nil {
			log.Panicln(err)
		}
		go func(acc net.Conn) {
			var read chan []byte
			var read2 chan []byte
			read = make(chan []byte, 1024)
			read2 = make(chan []byte, 1024)
			dl, err := net.Dial("tcp", remote)
			isClose := make(chan bool)
			isCloseDial := make(chan bool)
			if err != nil {
				log.Fatalln(err)
			}
			go func(ax net.Conn) {
				for {
					buf := make([]byte, 1024000)
					n, _ := ax.Read(buf)
					if len(buf[:n]) == 0 {
						isClose <- true
						return
					}
					read <- buf[:n]
				}
			}(acc)
			go func(ax2 net.Conn) {
				for {
					buf := make([]byte, 1024000)
					n, err := ax2.Read(buf)
					if len(buf[:n]) == 0 {
						if err != io.EOF {
							return
						}
						isCloseDial <- true
						return
					}
					read2 <- buf[:n]
				}
			}(dl)
			go func(dl net.Conn) {
				for {
					select {
					case b := <-read:
						if server {
							_, err := dl.Write(decryptAES(b))
							if err != nil {
								return
							}
						} else {
							_, err := dl.Write(encryptAES(b))
							if err != nil {
								return
							}
						}
						// _, err := dl.Write(RsaDecrypt(b))

					}
				}
			}(dl)
			defer func(c, v net.Conn) {
				_ = recover()
				log.Println("close")
				c.Close()
				v.Close()
				close(isCloseDial)
				close(isClose)
				close(read)
			}(acc, dl)
			for {
				select {
				case b := <-read2:
					// _, err := acc.Write(RsaEncrypt(b))

					if server {
						_, err := acc.Write(encryptAES(b))
						if err != nil {
							return
						}
					} else {
						_, err := acc.Write(decryptAES(b))
						if err != nil {
							return
						}
					}
				case <-isClose:
					panic("close")
				case <-isCloseDial:
					panic("close")
				}
			}

		}(a)
	}
}
